package lxs.swift.tool;


import lxs.swift.collector.sync.Mutex;
import org.jetbrains.annotations.NotNull;

import java.util.List;
import java.util.Random;

/**
 * 均衡器
 */
public class Balancer<T> {

    private final List<T> values;
    private final Mutex<Indexer> indexer;

    public Balancer(@NotNull List<T> values, @NotNull Indexer indexer) {
        if (values.size() != indexer.size()) throw new RuntimeException("size not match");
        this.values = values;
        this.indexer = new Mutex<>(indexer);
    }

    public T next() {
        return indexer.map(indexer -> {
            int i = indexer.nextIndex();
            return values.get(i);
        });
    }

    public interface Indexer {
        int nextIndex();

        int size();
    }

    public static class RandomIndexer implements Indexer {
        private final int size;
        private final Random random = new Random();

        public RandomIndexer(int size) {
            this.size = size;
        }

        @Override
        public int nextIndex() {
            return random.nextInt(size);
        }

        @Override
        public int size() {
            return size;
        }

    }

    public static class PollingIndexer implements Indexer {
        private final int size;
        private int index;

        public PollingIndexer(int size) {
            this.size = size;
            this.index = 0;
        }

        @Override
        public int nextIndex() {
            int i = index++;
            if (index >= size) index = 0;
            return i;
        }

        @Override
        public int size() {
            return size;
        }
    }

    public static class WeightedRandomIndexer implements Indexer {
        private final int[] weights;
        private final int sum;
        private final Random random = new Random();

        public WeightedRandomIndexer(int[] weights) {
            this.weights = weights;
            int sum = 0;
            for (int weight : weights) {
                sum += weight;
            }
            this.sum = sum;
        }

        @Override
        public int nextIndex() {
            int num = random.nextInt(sum);
            for (int i = 0; i < weights.length; i++) {
                int weight = weights[i];
                if (weight < num) return i;
                num -= weight;
            }
            return weights.length - 1;
        }

        @Override
        public int size() {
            return weights.length;
        }
    }

    public static class WeightedPollingIndexer implements Indexer {
        private final int[] weights;
        private final int sum;
        private int times;

        public WeightedPollingIndexer(int[] weights) {
            this.weights = weights;
            int sum = 0;
            for (int weight : weights) {
                sum += weight;
            }
            this.sum = sum;
            this.times = 0;
        }

        @Override
        public int nextIndex() {
            int num = (times++) % sum;
            for (int i = 0; i < weights.length; i++) {
                int weight = weights[i];
                if (weight < num) return i;
                num -= weight;
            }
            return weights.length - 1;
        }

        @Override
        public int size() {
            return weights.length;
        }
    }

    public static class SmoothWeightedPollingIndexer implements Indexer {
        private final int[] weights;
        private final int[] dynWeights;
        private final int sum;
        private int times;

        public SmoothWeightedPollingIndexer(int[] weights) {
            this.weights = weights;
            int sum = 0;
            for (int weight : weights) {
                sum += weight;
            }
            this.sum = sum;
            this.dynWeights = new int[weights.length];
        }

        @Override
        public int nextIndex() {
            int maxI = 0, maxWeight = 0;
            for (int i = 0; i < weights.length; i++) {
                int weight = weights[i] + dynWeights[i];
                if (i == 0) {
                    maxWeight = weight;
                    continue;
                }
                maxI = maxWeight > weight ? maxI : i;
            }
            dynWeights[maxI] -= sum;
            return maxI;
        }

        @Override
        public int size() {
            return weights.length;
        }
    }
}
